import pickle
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt

def unpickle(file):
    """解包CIFAR-10数据文件"""
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


class CIFAR10Dataset(Dataset):
    """自定义CIFAR-10数据集类"""

    def __init__(self, data_path, train=True, transform=None):
        self.transform = transform
        self.data = []
        self.labels = []

        if train:
            # 加载训练集
            for i in range(1, 6):
                file_path = os.path.join(data_path, f"data_batch_{i}")
                batch_dict = unpickle(file_path)
                self.data.append(batch_dict[b'data'])
                self.labels.extend(batch_dict[b'labels'])
        else:
            # 加载测试集
            file_path = os.path.join(data_path, "test_batch")
            batch_dict = unpickle(file_path)
            self.data.append(batch_dict[b'data'])
            self.labels.extend(batch_dict[b'labels'])

        # 合并所有数据
        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # 转换为HWC格式

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image = self.data[idx]
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label


def visualize_dataset(dataset, class_names, num_samples=20, title="CIFAR-10 Dataset Samples"):
    """
    可视化数据集中的样本图片
    """
    # 设置matplotlib中文字体（如果需要）
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    
    # 创建子图
    num_cols = 5
    num_rows = (num_samples + num_cols - 1) // num_cols
    fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 3 * num_rows))
    fig.suptitle(title, fontsize=16, fontweight='bold')
    
    # 如果只有一行，确保axes是二维数组
    if num_rows == 1:
        axes = axes.reshape(1, -1)
    
    # 随机选择样本
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    for i, idx in enumerate(indices):
        row = i // num_cols
        col = i % num_cols
        
        # 获取原始图像（不应用transform）
        image, label = dataset.data[idx], dataset.labels[idx]
        
        # 显示图像
        axes[row, col].imshow(image)
        axes[row, col].set_title(f'{class_names[label]} ({label})')
        axes[row, col].axis('off')
    
    # 隐藏多余的子图
    for i in range(num_samples, num_rows * num_cols):
        row = i // num_cols
        col = i % num_cols
        axes[row, col].axis('off')
    
    plt.tight_layout()
    plt.show()


def visualize_batch(data_loader, class_names, num_batches=1):
    """
    可视化数据加载器中的批次数据
    """
    data_iter = iter(data_loader)
    
    for batch_idx in range(num_batches):
        try:
            images, labels = next(data_iter)
        except StopIteration:
            break
            
        print(f"Batch {batch_idx + 1}:")
        print(f"Images shape: {images.shape}")
        print(f"Labels shape: {labels.shape}")
        print(f"Labels: {labels[:10].tolist()}...")  # 显示前10个标签
        
        # 显示批次中的前几个图像
        num_show = min(8, len(images))
        fig, axes = plt.subplots(2, 4, figsize=(15, 8))
        fig.suptitle(f'Batch {batch_idx + 1} - First {num_show} Images', fontsize=16)
        
        axes = axes.flatten()
        for i in range(num_show):
            # 将张量转换回numpy格式并调整通道顺序
            img = images[i].numpy().transpose(1, 2, 0)
            # 反标准化
            mean = np.array([0.4914, 0.4822, 0.4465])
            std = np.array([0.2023, 0.1994, 0.2010])
            img = std * img + mean
            img = np.clip(img, 0, 1)
            
            axes[i].imshow(img)
            axes[i].set_title(f'{class_names[labels[i].item()]} (Label: {labels[i].item()})')
            axes[i].axis('off')
        
        # 隐藏多余的子图
        for i in range(num_show, 8):
            axes[i].axis('off')
        
        plt.tight_layout()
        plt.show()


def get_data_loaders(batch_size=128, data_path='../cifar-10-batches-py'):
    """
    获取数据加载器
    """
    # 数据预处理
    transform_train = transforms.Compose([
        transforms.ToPILImage(),  # 将numpy数组转换为PIL图像
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    transform_test = transforms.Compose([
        transforms.ToPILImage(),  # 将numpy数组转换为PIL图像
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    # 加载数据集
    trainset = CIFAR10Dataset(data_path, train=True, transform=transform_train)
    testset = CIFAR10Dataset(data_path, train=False, transform=transform_test)

    # 获取类别名称
    meta_dict = unpickle(os.path.join(data_path, "batches.meta"))
    class_names = [label.decode('utf-8') for label in meta_dict[b'label_names']]

    # 创建数据加载器
    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
    testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)

    return trainloader, testloader, class_names, trainset, testset


def main():
    # 获取数据加载器
    data_path = '../cifar-10-batches-py'  # 根据你的实际路径修改
    trainloader, testloader, class_names, trainset, testset = get_data_loaders(data_path=data_path)

    print("CIFAR-10类别:", class_names)
    print(f"训练集大小: {len(trainset)}")
    print(f"测试集大小: {len(testset)}")

    # 可视化原始训练集样本
    print("\n可视化训练集原始样本:")
    visualize_dataset(trainset, class_names, num_samples=15, title="CIFAR-10 Training Set Samples")

    # 可视化原始测试集样本
    print("\n可视化测试集原始样本:")
    visualize_dataset(testset, class_names, num_samples=15, title="CIFAR-10 Test Set Samples")

    # 可视化经过预处理后的批次数据
    print("\n可视化预处理后的训练批次:")
    visualize_batch(trainloader, class_names, num_batches=1)

    # 获取一个批次的数据进行测试
    data_iter = iter(trainloader)
    images, labels = next(data_iter)
    print(f"\n图像批次形状: {images.shape}")
    print(f"标签批次形状: {labels.shape}")

    return trainloader, testloader, class_names


if __name__ == '__main__':
    # 如果是直接运行这个脚本，就执行main函数
    trainloader, testloader, class_names = main()

    # 如果你想在其他文件中导入这个模块，可以这样使用：
    # from CIFAR10Dataset import get_data_loaders
    # trainloader, testloader, class_names, trainset, testset = get_data_loaders()